import cv2
import numpy as np
from osgeo import gdal
import os,shutil
import glob
from super4XX.rrdbnet_arch import RRDBNet
from super4XX.realesrgan import RealESRGANer
import matplotlib.pyplot as plt

# 用于将单个tiff进行超分，会进行图像的裁剪和拼接

def file_name(file_dir):#返回某文件夹中后缀为png的所有文件绝对路径，为列表
    L=[]   
    for root, dirs, files in os.walk(file_dir):  
        for file in files:  
            if os.path.splitext(file)[1] == '.png':  
                L.append(os.path.join(root, file))  
    return L

def checkMkdir(path):
    if not os.path.exists(path):
        os.makedirs(path)

class Geotiff(object):
    '处理tiff影像文件，翻转平移缩放'
    def __init__(self,input_tif):
        self.input_tif = input_tif
        self.tiff_image = gdal.Open(self.input_tif,gdal.GA_Update)
        self.tiff_dir = os.path.dirname(input_tif)
        self.tiff_name = os.path.splitext(os.path.basename(input_tif))[0]
        self.geotransform = self.tiff_image.GetGeoTransform()
        self.coordorigin_X = self.geotransform[0]
        self.coordorigin_Y = self.geotransform[3]
        self.pixel_width = self.geotransform[1]
        self.pixel_height = self.geotransform[5]
        self.wtk = self.tiff_image.GetProjection()
        self.__get_upper_coefficient()

    @property
    def __geotransform(self):
        '''
        return:
        index [0 1 2 3 4 5]
        0-coefficients are zero
        [orignpoint_X,pixel width,coefficients,orign point_Y,coefficients,pixel height]
        '''
        return self.tiff_image.GetGeoTransform()

    @property
    def array(self):
        return self.tiff_image.ReadAsArray().astype(np.uint8)

    def __get_upper_coefficient(self):
        gt = self.geotransform
        self.upper_left = np.array([[gt[0]], [gt[3]]])
        self.coefficient = np.array([[gt[1], gt[2]],[gt[4], gt[5]]])
    
    def pixel2geo_trans(self,pixel_width,pixel_height):
        "像素坐标转换成地理坐标"
        gt = self.geotransform
        Xgeo = gt[0] + pixel_width*gt[1]  + pixel_height*gt[2]
        Ygeo = gt[3] + pixel_width*gt[4]  + pixel_height*gt[5]
        return [Xgeo,Ygeo]
    
    def pixels2geos_trans(self,pixel_coords):
        """[批量转换像素坐标到地理坐标]

        Args:
            geo_coords ([numpy array]): [像素标数组，任意维度但最后一维长度为2]

        Returns:
            [numpy array]: [转换后坐标成果]
        """
        shape = np.shape(pixel_coords)
        assert shape[-1] == 2, 'the last dimension of input pixel_coords is not equal 2'
        pixel_coords = np.reshape(pixel_coords.flatten(),(-1,2))
        # 'Affine transformation': W = A * pxl + ulnp.
        # world[2, 1] = a[2, 2] * pxl[2, 1] + upper_left[2,1]
        pixel_coords = np.transpose(pixel_coords)  # world coordinates
        geo_coords = np.matmul(self.coefficient, pixel_coords) + self.upper_left
        geo_coords = np.reshape(np.transpose(geo_coords),shape)
        return geo_coords
    
    def geo2pixel_trans(self,Xgeo,Ygeo):
        "地理坐标转换成像素坐标"
        # 'Affine transformation': W = A * pxl + ul
        # world[2, 1] = a[2, 2] * pxl[2, 1] + upper_left[2,1]
        world = np.array([[Xgeo], [Ygeo]])  # world coordinates
        # Reformulate: A^-1 * (W - ul) = pxl
        pxl = np.matmul(np.linalg.inv(self.coefficient), (world - self.upper_left))
        return pxl.flatten().tolist()
    
    def geos2pixels_trans(self,geo_coords):
        """[批量转换地理坐标到像素坐标]

        Args:
            geo_coords ([numpy array]): [地理坐标数组，任意维度但最后一维长度为2]

        Returns:
            [numpy array]: [转换后坐标成果]
        """
        shape = np.shape(geo_coords)
        assert shape[-1] == 2, 'the last dimension of input geo_coords is not equal 2'
        geo_coords = np.reshape(geo_coords.flatten(),(-1,2))
        # 'Affine transformation': W = A * pxl + ulnp.
        # world[2, 1] = a[2, 2] * pxl[2, 1] + upper_left[2,1]
        geo_coords = np.transpose(geo_coords)  # world coordinates
        pixel_coords = np.matmul(np.linalg.inv(self.coefficient), (geo_coords - self.upper_left))
        pixel_coords = np.reshape(np.transpose(pixel_coords),shape)
        return pixel_coords
    
    def geos2pixel(self,lst):
        new_lst = []
        for coords in lst:
            new_lst.append(self.geo2pixel_trans(*coords))
        return new_lst

class tiff2png(object):
    def __init__(self,tiff_path,pngindex):
        self.tiff_path = tiff_path
        self.save_dir = pngindex

    # 检查路径是否存在，不存在则创建
    def __checkMkdir(self):
        checkMkdir(self.save_dir)

    def origin2GT(self,image_array,name):
        #文件存储路径
        image_path = os.path.join(self.save_dir,"%s.png"%name)
        #imendcode方法可以解决中文路径乱码
        cv2.imencode('.png', image_array)[1].tofile(image_path)
            
    def create(self):
        self.__checkMkdir()
        print("decode file {0}".format(self.tiff_path))
        tiff_image = Geotiff(self.tiff_path)
        # if tiff_image.array.shape[0]<200 or tiff_image.array.shape[1]<10000:
        #     return 
        self.origin2GT(tiff_image.array, os.path.splitext(os.path.basename(self.tiff_path))[0])

def array2raster(outpath,array,oldtiff):
    ds=gdal.Open(oldtiff)
    geoTransform_old = ds.GetGeoTransform()
    im_proj      = ds.GetProjection()
    geoTransform = (geoTransform_old[0],geoTransform_old[1]/4,geoTransform_old[2],geoTransform_old[3],geoTransform_old[4],geoTransform_old[5]/4)
    cols=array.shape[1]
    rows=array.shape[0]

    driver=gdal.GetDriverByName('Gtiff')
    outRaster=driver.Create(outpath,cols,rows,1,gdal.GDT_Byte)
    outRaster.SetGeoTransform(geoTransform)#参数2,6为水平垂直分辨率，参数3,5表示图片是指北的
    outRaster.SetProjection(im_proj)
    outband=outRaster.GetRasterBand(1)
    outband.WriteArray(array)
    outRaster.FlushCache()
    del ds,outRaster

def super_resolution(model_path,input_png_filename,output_dir):
    
    #默认参数
    netscale = 4
    tile = 0
    tile_pad = 10
    pre_pad = 0
    half = False
    face_enhance = False
    ext = 'auto'
    outscale = 4
    block = 23

    if 'RealESRGAN_x4plus_anime_6B.pth' in model_path:
        block = 6
    elif 'RealESRGAN_x2plus.pth' in model_path:
        netscale = 2
    
    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=block, num_grow_ch=32, scale=netscale)

    upsampler = RealESRGANer(
        scale      = netscale,
        model_path = model_path,
        model      = model,
        tile       = tile,
        tile_pad   = tile_pad,
        pre_pad    = pre_pad,
        half       = half)

    if os.path.isfile(input_png_filename):
        paths = [input_png_filename]
    else:
        paths = sorted(glob.glob(os.path.join(input_png_filename, '*')))

    for idx, path in enumerate(paths):
        imgname, extension = os.path.splitext(os.path.basename(path))
        print('Testing', idx, imgname)

        img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
        if len(img.shape) == 3 and img.shape[2] == 4:
            img_mode = 'RGBA'
        else:
            img_mode = None

        try:
            if face_enhance:
                continue
            else:
                output, _ = upsampler.enhance(img, outscale=outscale)

        except Exception as error:
            print('Error', error)
            print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
        else:
            if ext == 'auto':
                extension = extension[1:]
            else:
                extension = ext
            if img_mode == 'RGBA':  # RGBA images should be saved in png format
                extension = 'png'

            save_path = os.path.join(output_dir, f'{imgname}.{extension}')
            cv2.imwrite(save_path, output)

def img_cut(big_png_filename,this_png_index):
    big_img = cv2.imread(big_png_filename,0)
    sss = 500
    cutnums = big_img.shape[1]//sss
    for i in range(cutnums):
        cutimg_name = this_png_index +'\\'+str(i)+'.png'
        cutimg = big_img[0:big_img.shape[0],i*sss:(i+1)*sss]
        cv2.imwrite(cutimg_name,cutimg)
    cutimg_name_fin = this_png_index +'\\'+str(cutnums)+'.png'
    cutimg = big_img[0:big_img.shape[0],cutnums*sss:big_img.shape[1]]
    cv2.imwrite(cutimg_name_fin,cutimg)

def img_buck(output_png_index):
    img = cv2.imread(output_png_index+'\\'+'0.png')
    for i in range(1,len(file_name(output_png_index))):
        img_plus = cv2.imread(output_png_index+'\\'+str(i)+'.png')
        img = np.hstack((img,img_plus))
    return img

def super_resolution_X4(filename,model_path):

    # 获取原始tiff名 无后缀
    tiff_name = os.path.splitext(os.path.basename(filename))[0]

    # 获取input png 文件夹名 大png
    input_png_dir = os.path.abspath(os.path.join(filename, "../..")+'\\'+'X1-PNG'+'\\'+tiff_name)
    checkMkdir(input_png_dir)

    # 裁切input png 文件夹名
    input_png_cut_dir = os.path.abspath(os.path.join(filename, "../..")+'\\'+'X1-PNG-CUT'+'\\'+tiff_name)
    checkMkdir(input_png_cut_dir)

    # 裁切output png 文件夹名
    output_png_cut_dir = os.path.abspath(os.path.join(filename, "../..")+'\\'+'X4-PNG-CUT'+'\\'+tiff_name)
    checkMkdir(output_png_cut_dir)

    # 获取output 4X tiff 文件夹名
    out_tif_dir = os.path.abspath(os.path.join(filename, "../.."))+'\\'+'X4-TIF'
    checkMkdir(out_tif_dir)

    # 取得input png 文件名
    input_png_name = input_png_dir + '\\'+ tiff_name + '.png'

    # 取得output tif 文件名
    output_tif_name = out_tif_dir + '\\'+ tiff_name + '.tif'

    print('start to change data type')
    # 将 geotiff 转为 png ,目标为同tif父级文件夹 X1-PNG
    tiffdata = tiff2png(filename,input_png_dir)
    tiffdata.create()
    print('done')

    print('start to cut image')
    # 先将大png裁切 放入input_png_cut_dir
    img_cut(input_png_name,input_png_cut_dir)
    print('done')

    print('start to super resolution')
    # 做超分
    cut_png_file_list = file_name(input_png_cut_dir)
    jj = 0
    for input_cut_png_name in cut_png_file_list :
        super_resolution(model_path,input_cut_png_name,output_png_cut_dir)
        jj = jj + 1 
        print('finish {0} in {1}'.format(jj,len(cut_png_file_list)))
    
    print('start to recover img')
    # 把小的拼起来
    img_biggg = img_buck(output_png_cut_dir)
    # 转tif
    array2raster(output_tif_name,img_biggg[:,:,0],filename)

    shutil.rmtree(input_png_dir)
    shutil.rmtree(input_png_cut_dir)
    shutil.rmtree(output_png_cut_dir)
    print('all finish')



